package controller

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/go-logr/logr"
	"go.uber.org/multierr"
	batchv1 "k8s.io/api/batch/v1"
	k8sapierror "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	"github.com/aquasecurity/trivy-operator/pkg/apis/aquasecurity/v1alpha1"
	"github.com/aquasecurity/trivy-operator/pkg/exposedsecretreport"
	"github.com/aquasecurity/trivy-operator/pkg/kube"
	"github.com/aquasecurity/trivy-operator/pkg/operator/etc"
	"github.com/aquasecurity/trivy-operator/pkg/sbomreport"
	"github.com/aquasecurity/trivy-operator/pkg/trivyoperator"
	"github.com/aquasecurity/trivy-operator/pkg/vulnerabilityreport"

	. "github.com/aquasecurity/trivy-operator/pkg/operator/predicate"
)

// streamReportToFile writes a report directly to file using streaming JSON encoding
// This reduces memory usage by avoiding intermediate marshaling to byte arrays
func streamReportToFile(report any, filePath string) error {
	file, err := os.Create(filePath)
	if err != nil {
		return fmt.Errorf("failed to create file %s: %w", filePath, err)
	}
	defer file.Close()

	encoder := json.NewEncoder(file)
	encoder.SetIndent("", "  ") // Pretty print for readability
	if err := encoder.Encode(report); err != nil {
		return fmt.Errorf("failed to encode report: %w", err)
	}

	return nil
}

// ScanJobController watches Kubernetes workloads and generates
// v1alpha1.VulnerabilityReport instances using vulnerability scanner that that
// implements the Plugin interface.
type ScanJobController struct {
	logr.Logger
	etc.Config
	kube.ObjectResolver
	kube.LogsReader
	vulnerabilityreport.Plugin
	trivyoperator.PluginContext
	trivyoperator.ConfigData
	SbomReadWriter          sbomreport.ReadWriter
	VulnerabilityReadWriter vulnerabilityreport.ReadWriter
	ExposedSecretReadWriter exposedsecretreport.ReadWriter
}

// Manage scan jobs with image pull secrets
// kubebuilder:rbac:groups="",resources=secrets,verbs=create;update
// +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;delete

func (r *ScanJobController) SetupWithManager(mgr ctrl.Manager) error {
	var predicates []predicate.Predicate
	if !r.ConfigData.VulnerabilityScanJobsInSameNamespace() {
		predicates = append(predicates, InNamespace(r.Config.Namespace))
	}
	predicates = append(predicates, ManagedByTrivyOperator, IsVulnerabilityReportScan, JobHasAnyCondition)
	return ctrl.NewControllerManagedBy(mgr).
		For(&batchv1.Job{}, builder.WithPredicates(predicates...)).
		Complete(r.reconcileJobs())
}

func (r *ScanJobController) reconcileJobs() reconcile.Func {
	return func(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
		log := r.Logger.WithValues("job", req.NamespacedName)

		job := &batchv1.Job{}
		err := r.Client.Get(ctx, req.NamespacedName, job)
		if err != nil {
			if k8sapierror.IsNotFound(err) {
				log.V(1).Info("Ignoring cached job that must have been deleted")
				return ctrl.Result{}, nil
			}
			return ctrl.Result{}, fmt.Errorf("getting job from cache: %w", err)
		}

		if len(job.Status.Conditions) == 0 {
			log.V(1).Info("Ignoring Job without conditions")
			return ctrl.Result{}, nil
		}

		switch jobCondition := job.Status.Conditions[0].Type; jobCondition {
		case batchv1.JobComplete, batchv1.JobSuccessCriteriaMet, batchv1.JobFailed, batchv1.JobFailureTarget:
			completedContainers, err := r.completedContainers(ctx, job)
			if err != nil {
				return ctrl.Result{}, r.deleteJob(ctx, job)
			}
			if len(completedContainers) == 0 {
				return ctrl.Result{}, r.deleteJob(ctx, job)
			}
			return ctrl.Result{}, r.processCompleteScanJob(ctx, job, completedContainers...)

		default:
			return ctrl.Result{}, fmt.Errorf("unrecognized scan job condition: %v", jobCondition)
		}
	}
}

func (r *ScanJobController) processCompleteScanJob(ctx context.Context, job *batchv1.Job, completedContainers ...string) error {
	log := r.Logger.WithValues("job", fmt.Sprintf("%s/%s", job.Namespace, job.Name))

	ownerRef, err := kube.ObjectRefFromObjectMeta(job.ObjectMeta)
	if err != nil {
		return fmt.Errorf("getting owner ref from scan job metadata: %w", err)
	}

	owner, err := r.ObjectFromObjectRef(ctx, ownerRef)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			log.V(1).Info("Report owner must have been deleted", "owner", owner)
			return r.deleteJob(ctx, job)
		}
		return fmt.Errorf("getting object from object ref: %w", err)
	}
	podSpecHash, ok := job.Labels[trivyoperator.LabelResourceSpecHash]
	if !ok {
		return fmt.Errorf("expected label %s not set", trivyoperator.LabelResourceSpecHash)
	}

	log = log.WithValues("kind", owner.GetObjectKind().GroupVersionKind().Kind,
		"name", owner.GetName(), "namespace", owner.GetNamespace(), "podSpecHash", podSpecHash)

	log.V(1).Info("Job complete")

	hasVulnReports := true
	containerImages, err := kube.GetContainerImagesFromJob(job, completedContainers...)
	if err != nil {
		return fmt.Errorf("getting container images: %w", err)
	}
	if r.Config.VulnerabilityScannerEnabled {
		hasVulnReports, err = hasVulnerabilityReports(ctx, r.VulnerabilityReadWriter, ownerRef, podSpecHash, containerImages)
		if err != nil {
			return err
		}
	}
	hasExposedSecretReports := true
	if r.Config.ExposedSecretScannerEnabled {
		hasExposedSecretReports, err = hasSecretReports(ctx, r.ExposedSecretReadWriter, ownerRef, podSpecHash, containerImages)
		if err != nil {
			return err
		}
	}
	if hasVulnReports && hasExposedSecretReports {
		log.V(1).Info("VulnerabilityReports already exist", "owner", owner)
		log.V(1).Info("Deleting complete scan job", "owner", owner)
		return r.deleteJob(ctx, job)
	}

	var vulnerabilityReports []v1alpha1.VulnerabilityReport
	var clusterVulnerabilityReports []v1alpha1.ClusterVulnerabilityReport
	var secretReports []v1alpha1.ExposedSecretReport
	var sbomNameSpacedReports []v1alpha1.SbomReport
	var sbomClusterReports []v1alpha1.ClusterSbomReport

	var merr error
	for containerName, containerImage := range containerImages {
		vulnReports, secReports, sbomReports, err := r.processScanJobResults(ctx, job, containerName, containerImage, owner)
		if err != nil {
			merr = multierr.Append(merr, err)
			continue
		}
		if vulnReports.vulnerabilityNamespaceReports != nil {
			vulnerabilityReports = append(vulnerabilityReports, *vulnReports.vulnerabilityNamespaceReports)
		}
		if vulnReports.vulnerabilityClusterReports != nil {
			clusterVulnerabilityReports = append(clusterVulnerabilityReports, *vulnReports.vulnerabilityClusterReports)
		}
		secretReports = append(secretReports, secReports...)
		if sbomReports != nil {
			sbomNameSpacedReports = append(sbomNameSpacedReports, sbomReports.sbomNamespaceReports...)
			sbomClusterReports = append(sbomClusterReports, sbomReports.sbomClusterReports...)
		}
	}
	if merr != nil {
		return merr
	}

	if r.Config.VulnerabilityScannerEnabled {
		if len(vulnerabilityReports) > 0 {
			err = r.VulnerabilityReadWriter.Write(ctx, vulnerabilityReports)
			if err != nil {
				return err
			}
		}
		if len(clusterVulnerabilityReports) > 0 {
			err = r.VulnerabilityReadWriter.WriteCluster(ctx, clusterVulnerabilityReports)
			if err != nil {
				return err
			}
		}
	}

	if r.Config.ExposedSecretScannerEnabled {
		err = r.ExposedSecretReadWriter.Write(ctx, secretReports)
		if err != nil {
			return err
		}
	}

	if r.Config.SbomGenerationEnable {
		err = r.SbomReadWriter.Write(ctx, sbomNameSpacedReports)
		if err != nil {
			return err
		}
	}
	if r.Config.ClusterSbomCacheEnable {
		if !r.Config.AltReportStorageEnabled || r.Config.AltReportDir == "" {
			err = r.SbomReadWriter.WriteCluster(ctx, sbomClusterReports)
			if err != nil {
				return err
			}
		}
	}

	log.V(1).Info("Deleting complete scan job", "owner", owner)
	return r.deleteJob(ctx, job)
}

type SbomReports struct {
	sbomNamespaceReports []v1alpha1.SbomReport
	sbomClusterReports   []v1alpha1.ClusterSbomReport
}

type VulnerabilityReports struct {
	vulnerabilityNamespaceReports *v1alpha1.VulnerabilityReport
	vulnerabilityClusterReports   *v1alpha1.ClusterVulnerabilityReport
}

func (r *ScanJobController) processScanJobResults(ctx context.Context,
	job *batchv1.Job,
	containerName,
	containerImage string,
	owner client.Object) (VulnerabilityReports, []v1alpha1.ExposedSecretReport, *SbomReports, error) {
	log := r.Logger.WithValues("job-results-processor", fmt.Sprintf("%s/%s", job.Namespace, job.Name))

	var vulnerabilityReports VulnerabilityReports
	var secretReports []v1alpha1.ExposedSecretReport
	sbomReports := &SbomReports{}

	podSpecHash, ok := job.Labels[trivyoperator.LabelResourceSpecHash]
	if !ok {
		return VulnerabilityReports{}, nil, nil, fmt.Errorf("expected label %s not set", trivyoperator.LabelResourceSpecHash)
	}

	logsStream, err := r.LogsReader.GetLogsByJobAndContainerName(ctx, job, containerName)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			log.V(1).Info("Cached job must have been deleted")
			return VulnerabilityReports{}, nil, nil, nil
		}
		if kube.IsPodControlledByJobNotFound(err) {
			log.V(1).Info("Pod must have been deleted")
			return VulnerabilityReports{}, nil, nil, r.deleteJob(ctx, job)
		}
		return VulnerabilityReports{}, nil, nil, fmt.Errorf("getting logs for pod %q: %w", job.Namespace+"/"+job.Name, err)
	}

	defer func() {
		err := logsStream.Close()
		if err != nil {
			log.V(1).Error(err, "could not close log stream")
		}
	}()

	vulnReportData, secretReportData, sbomReportData, err := r.Plugin.ParseReportData(r.PluginContext, containerImage, logsStream)
	if err != nil {
		return VulnerabilityReports{}, nil, nil, err
	}

	resourceLabelsToInclude := r.GetReportResourceLabels()
	additionalCustomLabels, err := r.GetAdditionalReportLabels()
	if err != nil {
		return VulnerabilityReports{}, nil, nil, err
	}
	if r.Config.VulnerabilityScannerEnabled {
		reportBuilder := vulnerabilityreport.NewReportBuilder(r.Client.Scheme()).
			Controller(owner).
			Container(containerName).
			Data(vulnReportData).
			PodSpecHash(podSpecHash).
			ResourceLabelsToInclude(resourceLabelsToInclude).
			AdditionalReportLabels(additionalCustomLabels)

		if r.Config.ScannerReportTTL != nil {
			reportBuilder.ReportTTL(r.Config.ScannerReportTTL)
		}

		report, clusterReport, err := reportBuilder.Get()
		if err != nil {
			return VulnerabilityReports{}, nil, nil, err
		}
		vulnerabilityReports = VulnerabilityReports{vulnerabilityNamespaceReports: report, vulnerabilityClusterReports: clusterReport}
	}
	_, reused := job.Labels[trivyoperator.LabelReusedReport]
	if !ok {
		return VulnerabilityReports{}, nil, nil, fmt.Errorf("expected label %s not set", trivyoperator.LabelResourceSpecHash)
	}

	if r.ExposedSecretScannerEnabled && !reused {
		secretReport, err := exposedsecretreport.NewReportBuilder(r.Client.Scheme()).
			Controller(owner).
			Container(containerName).
			Data(secretReportData).
			PodSpecHash(podSpecHash).
			ResourceLabelsToInclude(resourceLabelsToInclude).
			AdditionalReportLabels(additionalCustomLabels).
			Get()
		if err != nil {
			return VulnerabilityReports{}, nil, nil, err
		}
		secretReports = append(secretReports, secretReport)
	}

	if r.SbomGenerationEnable && sbomReportData != nil && !reused {
		sbomReportBuilder := sbomreport.NewReportBuilder(r.Client.Scheme()).
			Controller(owner).
			Container(containerName).
			Data(*sbomReportData).
			PodSpecHash(podSpecHash).
			CacheTTL(r.Config.CacheReportTTL).
			ResourceLabelsToInclude(resourceLabelsToInclude).
			AdditionalReportLabels(additionalCustomLabels)
		sbomReport, clusterReport, err := sbomReportBuilder.Get()
		if err != nil {
			return VulnerabilityReports{}, nil, nil, err
		}
		if r.Config.ClusterSbomCacheEnable {
			sbomReports.sbomClusterReports = []v1alpha1.ClusterSbomReport{clusterReport}
		}
		sbomReports.sbomNamespaceReports = []v1alpha1.SbomReport{sbomReport}
	}
	if r.Config.AltReportStorageEnabled && r.Config.AltReportDir != "" {
		log.V(1).Info("Writing vulnerability reports to alternate storage", "dir", r.Config.AltReportDir)
		// Initilizing alternate writing directory var for trivy reports
		reportDir := r.Config.AltReportDir
		// Create subdirectories for each type of report
		clusterVulnerabilityDir := filepath.Join(reportDir, "cluster_vulnerability_reports")
		vulnerabilityDir := filepath.Join(reportDir, "vulnerability_reports")
		secretDir := filepath.Join(reportDir, "secret_reports")
		clusterSbomDir := filepath.Join(reportDir, "cluster_sbom_reports")
		sbomDir := filepath.Join(reportDir, "sbom_reports")

		// Ensure the directories exist
		for _, dir := range []string{vulnerabilityDir, secretDir, sbomDir, clusterSbomDir, clusterVulnerabilityDir} {
			if err := os.MkdirAll(dir, 0o750); err != nil {
				return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, fmt.Errorf("failed to make directory %s: %w", dir, err)
			}
		}

		// Get workload kind and name
		workloadKind := owner.GetObjectKind().GroupVersionKind().Kind
		workloadName := owner.GetName()
		// Write cluster vulnerability reports to a file using streaming
		if vulnerabilityReports.vulnerabilityClusterReports != nil {
			reportPath := filepath.Join(clusterVulnerabilityDir, fmt.Sprintf("%s-%s-%s.json", workloadKind, workloadName, containerName))
			if err := streamReportToFile(vulnerabilityReports.vulnerabilityClusterReports, reportPath); err != nil {
				return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, fmt.Errorf("failed to write cluster vulnerability report: %w", err)
			}
		}

		// Write vulnerability reports to a file using streaming
		if vulnerabilityReports.vulnerabilityNamespaceReports != nil {
			reportPath := filepath.Join(vulnerabilityDir, fmt.Sprintf("%s-%s-%s.json", workloadKind, workloadName, containerName))
			if err := streamReportToFile(vulnerabilityReports.vulnerabilityNamespaceReports, reportPath); err != nil {
				return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, fmt.Errorf("failed to write vulnerability report: %w", err)
			}
		}

		// Write secret reports to a file using streaming
		if len(secretReports) > 0 {
			reportPath := filepath.Join(secretDir, fmt.Sprintf("%s-%s-%s.json", workloadKind, workloadName, containerName))
			if err := streamReportToFile(secretReports, reportPath); err != nil {
				return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, fmt.Errorf("failed to write exposed secrets report: %w", err)
			}
		}

		// Write SBOM cluster reports to a file using streaming
		if len(sbomReports.sbomClusterReports) > 0 {
			reportPath := filepath.Join(clusterSbomDir, fmt.Sprintf("%s-%s-%s.json", workloadKind, workloadName, containerName))
			if err := streamReportToFile(sbomReports.sbomClusterReports, reportPath); err != nil {
				return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, fmt.Errorf("failed to write sbom cluster report: %w", err)
			}
		}

		// Write SBOM reports to a file using streaming
		if len(sbomReports.sbomNamespaceReports) > 0 {
			reportPath := filepath.Join(sbomDir, fmt.Sprintf("%s-%s-%s.json", workloadKind, workloadName, containerName))
			if err := streamReportToFile(sbomReports.sbomNamespaceReports, reportPath); err != nil {
				return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, fmt.Errorf("failed to write sbom report: %w", err)
			}
		}
		// Return empty reports since alternate write method is used
		return VulnerabilityReports{}, []v1alpha1.ExposedSecretReport{}, &SbomReports{}, nil
	}
	return vulnerabilityReports, secretReports, sbomReports, nil
}

func (r *ScanJobController) completedContainers(ctx context.Context, scanJob *batchv1.Job) ([]string, error) {
	log := r.Logger.WithValues("job", fmt.Sprintf("%s/%s", scanJob.Namespace, scanJob.Name))

	statuses, err := r.GetTerminatedContainersStatusesByJob(ctx, scanJob)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			log.V(1).Info("Cached job must have been deleted")
			return []string{}, nil
		}
		if kube.IsPodControlledByJobNotFound(err) {
			log.V(1).Info("Pod must have been deleted")
			return []string{}, nil
		}
		return nil, err
	}
	completedContainers := make([]string, 0)
	for container, status := range statuses {
		if status.ExitCode == 0 {
			completedContainers = append(completedContainers, container)
			continue
		}
		if strings.Contains(status.Message, "no child with platform linux") {
			log.Info("Scan job container", "container", container, "status.reason", status.Reason, "status.message", "By default, only Linux amd64 images are supported for scanning.")
		} else {
			log.Error(nil, "Scan job container", "container", container, "status.reason", status.Reason, "status.message", status.Message)
		}
	}
	return completedContainers, nil
}

func (r *ScanJobController) deleteJob(ctx context.Context, job *batchv1.Job) error {
	if job.Spec.TTLSecondsAfterFinished != nil {
		return nil
	}

	err := r.Client.Delete(ctx, job, client.PropagationPolicy(metav1.DeletePropagationBackground))
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			return nil
		}
		return fmt.Errorf("deleting job: %w", err)
	}
	return nil
}
